#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时脑电控制机器狗演示
"""

import numpy as np
import time
import threading
import tkinter as tk
from tkinter import ttk
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from eeg_dog_control_model import RealTimeEEGController, DogController, LightweightEEGControlModel

class EEGSimulator:
    """
    脑电信号模拟器
    用于演示实时控制
    """
    
    def __init__(self, sampling_rate=1000):
        self.sampling_rate = sampling_rate
        self.current_pattern = 0  # 当前模式
        self.time_counter = 0
        
    def generate_eeg_signal(self, pattern_type=0):
        """
        生成模拟脑电信号
        """
        # 基础信号
        t = self.time_counter / self.sampling_rate
        base_signal = np.random.randn() * 0.1
        
        # 根据模式添加特定成分
        if pattern_type == 0:  # 向前 - 高频
            pattern = np.sin(2 * np.pi * 15 * t) * 0.3
        elif pattern_type == 1:  # 向后 - 低频
            pattern = np.sin(2 * np.pi * 8 * t) * 0.3
        elif pattern_type == 2:  # 趴下 - 中频
            pattern = np.sin(2 * np.pi * 12 * t) * 0.3
        else:  # 转圈 - 复杂模式
            pattern = np.sin(2 * np.pi * 10 * t) * 0.3 + np.sin(2 * np.pi * 5 * t) * 0.2
        
        signal = base_signal + pattern
        self.time_counter += 1
        
        return signal

class RealTimeControlGUI:
    """
    实时控制GUI界面
    """
    
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("脑电控制机器狗演示")
        self.root.geometry("800x600")
        
        # 初始化组件
        self.setup_gui()
        
        # 初始化控制器
        self.setup_controllers()
        
        # 启动实时处理
        self.start_realtime_processing()
        
    def setup_gui(self):
        """
        设置GUI界面
        """
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 控制面板
        control_frame = ttk.LabelFrame(main_frame, text="控制面板", padding="10")
        control_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        
        # 模式选择
        ttk.Label(control_frame, text="选择脑电模式:").grid(row=0, column=0, sticky=tk.W)
        self.pattern_var = tk.IntVar(value=0)
        pattern_names = ['向前', '向后', '趴下', '转圈']
        for i, name in enumerate(pattern_names):
            ttk.Radiobutton(control_frame, text=name, variable=self.pattern_var, 
                           value=i, command=self.on_pattern_change).grid(row=0, column=i+1, padx=5)
        
        # 状态显示
        status_frame = ttk.LabelFrame(main_frame, text="状态信息", padding="10")
        status_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        
        self.status_label = ttk.Label(status_frame, text="等待连接...")
        self.status_label.grid(row=0, column=0, sticky=tk.W)
        
        self.confidence_label = ttk.Label(status_frame, text="置信度: 0.000")
        self.confidence_label.grid(row=1, column=0, sticky=tk.W)
        
        self.action_label = ttk.Label(status_frame, text="当前动作: 无")
        self.action_label.grid(row=2, column=0, sticky=tk.W)
        
        # 信号显示
        signal_frame = ttk.LabelFrame(main_frame, text="脑电信号", padding="10")
        signal_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        
        # 创建matplotlib图形
        self.fig, self.ax = plt.subplots(figsize=(8, 3))
        self.canvas = FigureCanvasTkAgg(self.fig, signal_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # 按钮面板
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=3, column=0, columnspan=2, pady=(10, 0))
        
        ttk.Button(button_frame, text="开始演示", command=self.start_demo).grid(row=0, column=0, padx=5)
        ttk.Button(button_frame, text="停止演示", command=self.stop_demo).grid(row=0, column=1, padx=5)
        ttk.Button(button_frame, text="紧急停止", command=self.emergency_stop).grid(row=0, column=2, padx=5)
        
        # 配置网格权重
        main_frame.columnconfigure(0, weight=1)
        main_frame.rowconfigure(2, weight=1)
        
    def setup_controllers(self):
        """
        设置控制器
        """
        # 创建模拟模型
        self.model = LightweightEEGControlModel()
        
        # 创建机器狗控制器
        self.dog_controller = DogController()
        
        # 创建脑电模拟器
        self.eeg_simulator = EEGSimulator()
        
        # 创建实时控制器
        self.eeg_controller = RealTimeEEGController(
            model_path=None,  # 使用内存中的模型
            dog_controller=self.dog_controller
        )
        self.eeg_controller.model = self.model
        
        # 初始化数据缓冲区
        self.signal_buffer = []
        self.max_buffer_size = 500
        
        # 控制标志
        self.is_running = False
        self.demo_thread = None
        
    def on_pattern_change(self):
        """
        模式改变回调
        """
        pattern = self.pattern_var.get()
        pattern_names = ['向前', '向后', '趴下', '转圈']
        self.status_label.config(text=f"当前模式: {pattern_names[pattern]}")
        
    def start_demo(self):
        """
        开始演示
        """
        if not self.is_running:
            self.is_running = True
            self.demo_thread = threading.Thread(target=self.demo_loop)
            self.demo_thread.daemon = True
            self.demo_thread.start()
            print("🚀 开始实时控制演示")
    
    def stop_demo(self):
        """
        停止演示
        """
        self.is_running = False
        print("⏹️ 停止演示")
    
    def emergency_stop(self):
        """
        紧急停止
        """
        self.is_running = False
        print("🚨 紧急停止！")
        # 这里应该发送停止指令给机器狗
    
    def demo_loop(self):
        """
        演示循环
        """
        while self.is_running:
            try:
                # 生成脑电信号
                pattern = self.pattern_var.get()
                eeg_sample = self.eeg_simulator.generate_eeg_signal(pattern)
                
                # 添加到缓冲区
                self.signal_buffer.append(eeg_sample)
                if len(self.signal_buffer) > self.max_buffer_size:
                    self.signal_buffer.pop(0)
                
                # 当缓冲区满时进行预测
                if len(self.signal_buffer) == self.max_buffer_size:
                    self.process_prediction()
                
                # 更新GUI
                self.update_signal_display()
                
                # 控制更新频率
                time.sleep(0.01)  # 100Hz更新频率
                
            except Exception as e:
                print(f"❌ 演示循环出错: {e}")
                break
    
    def process_prediction(self):
        """
        处理预测
        """
        try:
            # 准备输入数据
            input_data = np.array(self.signal_buffer).reshape(1, 5, 100, 1, 1)
            
            # 模型预测
            prediction = self.model.predict(input_data, verbose=0)
            control_class = np.argmax(prediction[0])
            confidence = np.max(prediction[0])
            
            # 更新GUI
            self.root.after(0, self.update_prediction_display, control_class, confidence)
            
            # 执行控制（如果置信度足够高）
            if confidence > 0.7:
                self.root.after(0, self.execute_control, control_class)
                
        except Exception as e:
            print(f"❌ 预测处理出错: {e}")
    
    def update_prediction_display(self, control_class, confidence):
        """
        更新预测显示
        """
        pattern_names = ['向前', '向后', '趴下', '转圈']
        self.confidence_label.config(text=f"置信度: {confidence:.3f}")
        self.action_label.config(text=f"预测动作: {pattern_names[control_class]}")
    
    def execute_control(self, control_class):
        """
        执行控制
        """
        control_actions = {
            0: self.dog_controller.move_forward,
            1: self.dog_controller.move_backward,
            2: self.dog_controller.lie_down,
            3: self.dog_controller.turn_around
        }
        
        if control_class in control_actions:
            action_name = ['向前', '向后', '趴下', '转圈'][control_class]
            print(f"🎯 执行控制: {action_name}")
            control_actions[control_class]()
    
    def update_signal_display(self):
        """
        更新信号显示
        """
        try:
            self.ax.clear()
            
            # 显示最近的信号
            if len(self.signal_buffer) > 0:
                signal_data = self.signal_buffer[-100:]  # 显示最近100个点
                time_axis = np.arange(len(signal_data)) / 1000  # 转换为秒
                
                self.ax.plot(time_axis, signal_data, 'b-', linewidth=1)
                self.ax.set_xlabel('时间 (秒)')
                self.ax.set_ylabel('振幅')
                self.ax.set_title('实时脑电信号')
                self.ax.grid(True, alpha=0.3)
                
                # 设置y轴范围
                self.ax.set_ylim(-1, 1)
            
            self.canvas.draw()
            
        except Exception as e:
            print(f"❌ 信号显示更新出错: {e}")
    
    def start_realtime_processing(self):
        """
        启动实时处理
        """
        print("🎮 实时控制界面启动")
        self.root.mainloop()

def main():
    """
    主函数
    """
    print("🧠 脑电控制机器狗实时演示")
    print("=" * 50)
    print("使用说明:")
    print("1. 选择脑电模式（向前/向后/趴下/转圈）")
    print("2. 点击'开始演示'启动实时控制")
    print("3. 观察信号波形和预测结果")
    print("4. 点击'停止演示'或'紧急停止'")
    print("=" * 50)
    
    # 启动GUI
    app = RealTimeControlGUI()

if __name__ == "__main__":
    main() 